package container

import (
	"sync"

	"golang.org/x/sync/singleflight"
)

type MapInterface[K comparable, V any] interface {
	Load(key K) (V, bool)
	Store(key K, value V)
	LoadOrStore(key K, value V) (actual V, loaded bool)
	LoadOrStoreNew(key K, ctor func() V, release func(V), keyString string) (actual V, loaded bool)
	LoadOrStoreNewWithError(key K, ctor func() (V, error), release func(V), keyString string) (actual V, loaded bool, err error)
	LoadAndDelete(key K) (value V, loaded bool)
	Delete(key K)
	Range(func(key K, value V) (shouldContinue bool))
	Swap(key K, value V) (previous V, loaded bool)
	CompareAndSwap(key K, old V, new V) (swapped bool)
	CompareAndDelete(key K, old V) (deleted bool)
}

type syncMap[K comparable, V any] struct {
	it           sync.Map
	sf           singleflight.Group
	defaultValue V
}

func NewSyncMap[K comparable, V any]() *syncMap[K, V] {
	var defaultValue V
	return &syncMap[K, V]{
		defaultValue: defaultValue,
	}
}

func (s *syncMap[K, V]) Load(key K) (V, bool) {
	v, ok := s.it.Load(key)
	if !ok {
		return s.defaultValue, false
	}
	return v.(V), true
}

func (s *syncMap[K, V]) Store(key K, value V) {
	s.it.Store(key, value)
}

func (s *syncMap[K, V]) LoadOrStore(key K, value V) (actual V, loaded bool) {
	v, l := s.it.LoadOrStore(key, value)
	return v.(V), l
}

func (s *syncMap[K, V]) LoadAndDelete(key K) (value V, loaded bool) {
	v, l := s.it.LoadAndDelete(key)
	if !l {
		return s.defaultValue, false
	}
	return v.(V), true
}

func (s *syncMap[K, V]) Delete(key K) {
	s.it.Delete(key)
}

func (s *syncMap[K, V]) Range(f func(key K, value V) (shouldContinue bool)) {
	s.it.Range(func(key, value any) bool {
		return f(key.(K), value.(V))
	})
}

func (s *syncMap[K, V]) Swap(key K, value V) (previous V, loaded bool) {
	p, l := s.it.Swap(key, value)
	if !l {
		return s.defaultValue, false
	}
	return p.(V), true
}

func (s *syncMap[K, V]) CompareAndSwap(key K, old V, new V) (swapped bool) {
	return s.it.CompareAndSwap(key, old, new)
}

func (s *syncMap[K, V]) CompareAndDelete(key K, old V) (deleted bool) {
	return s.it.CompareAndDelete(key, old)
}

func (s *syncMap[K, V]) LoadOrStoreNew(key K, ctor func() V, release func(V), keyString string) (actual V, loaded bool) {
	v, l := s.it.Load(key)
	if l {
		return v.(V), l
	}
	ret, _, _ := s.sf.Do(keyString, func() (interface{}, error) {
		newValue := ctor()
		tempV, tempL := s.LoadOrStore(key, newValue)
		if tempL && release != nil {
			release(newValue)
		}
		return loadStatus[V]{
			v: tempV,
			l: tempL,
		}, nil
	})
	loadVStatus := ret.(loadStatus[V])
	return loadVStatus.v, loadVStatus.l
}

func (s *syncMap[K, V]) LoadOrStoreNewWithError(key K, ctor func() (V, error), release func(V), keyString string) (actual V, loaded bool, err error) {
	var nilV V
	v, l := s.it.Load(key)
	if l {
		return v.(V), l, nil
	}
	ret, err, _ := s.sf.Do(keyString, func() (interface{}, error) {
		newValue, err := ctor()
		if err != nil {
			return nil, err
		}
		tempV, tempL := s.LoadOrStore(key, newValue)
		if tempL && release != nil {
			release(newValue)
		}
		return loadStatus[V]{
			v: tempV,
			l: tempL,
		}, nil
	})
	if err != nil {
		return nilV, false, err
	}
	loadVStatus := ret.(loadStatus[V])
	return loadVStatus.v, loadVStatus.l, nil
}

type loadStatus[V any] struct {
	v V
	l bool
}
